package com.jay.x.net.ssl;

import java.io.FileInputStream;
import java.io.IOException;
import java.net.InetAddress;
import java.net.InetSocketAddress;
import java.nio.ByteBuffer;
import java.nio.channels.SelectionKey;
import java.nio.channels.Selector;
import java.nio.channels.ServerSocketChannel;
import java.nio.channels.SocketChannel;
import java.security.KeyStore;
import java.util.Iterator;
import java.util.Set;

import javax.net.ssl.KeyManagerFactory;
import javax.net.ssl.SSLContext;
import javax.net.ssl.SSLEngine;

public class SSLNIOEchoServer {
	SSLContext context;
	ServerSocketChannel ssc;
	Selector sel;

	public SSLNIOEchoServer() throws Exception {
		// Create the SSLContext
		this.context = SSLContext.getInstance("TLS");
		// Initialize KMF ...
		KeyManagerFactory kmf = KeyManagerFactory.getInstance("SunX509");
		KeyStore ks = KeyStore.getInstance("JKS");
		char[] password = System.getProperty("javax.net.ssl.keyStorePassword")
				.toCharArray();
		ks.load(new FileInputStream(System
				.getProperty("javax.net.ssl.keyStore")), password);
		kmf.init(ks, password);
		context.init(kmf.getKeyManagers(), null, null);
		// Start the server
		this.ssc = ServerSocketChannel.open();
		ssc.configureBlocking(false);
		ssc.socket().bind(new InetSocketAddress((InetAddress) null, 0), 50);
		System.out.println("Server: listening at " + ssc);
		this.sel = Selector.open();
		ssc.register(sel, SelectionKey.OP_ACCEPT);
	}

	public void run() throws IOException {
		// Selector loop
		int count;
		while (sel.keys().size() > 0) {
			try {
				count = sel.select(30 * 1000);
				if (count < 0) {
					System.out.println("Server: select timeout");
					continue;
				}
			} catch (IOException exc) {
				exc.printStackTrace();
				sel.close();
				ssc.close();
				return;
			}
			System.out.println("Server: select count=" + count);
			Set<SelectionKey> selKeys = sel.selectedKeys();
			Iterator<SelectionKey> it = selKeys.iterator();
			// process ready keys
			while (it.hasNext()) {
				SelectionKey sk = (SelectionKey) it.next();
				it.remove();
				if (!sk.isValid())
					continue;
				try {
					if (sk.isAcceptable())
						handleAccept(sk);
					if (sk.isReadable())
						handleRead(sk);
					if (sk.isWritable())
						handleWrite(sk);
				} catch (IOException exc) {
					exc.printStackTrace();
					sk.channel().close();
				}
			}
		}
	} // run()
		// Process OP_ACCEPT

	void handleAccept(SelectionKey sk) throws IOException {
		ServerSocketChannel ssc = (ServerSocketChannel) sk.channel();
		SocketChannel sc = ssc.accept();
		if (sc != null) {
			System.out.println("Server: accepted " + sc);
			sc.configureBlocking(false);
			// Create an SSL engine for this connection
			SSLEngine engine = context.createSSLEngine("localhost", sc.socket()
					.getPort());
			// This is the server end
			engine.setUseClientMode(false);
			// Create engine manager for the channel & engine
			SSLEngineManager mgr = new SSLEngineManager(sc, engine);
			// Register for OP_READ with mgr as attachment
			sc.register(sel, SelectionKey.OP_READ, mgr);
		}
	}

	// process OP_READ
	void handleRead(SelectionKey sk) throws IOException {
		SSLEngineManager mgr = (SSLEngineManager) sk.attachment();
		SSLEngine engine = mgr.getEngine();
		ByteBuffer request = mgr.getAppRecvBuffer();
		System.out.println("Server: reading");
		int count = mgr.read();
		System.out.println("Server: read count=" + count + " request="
				+ request);
		if (count < 0) {
			// client has closed
			mgr.close();
			// finished with this key
			sk.cancel();
			// finished with this test actually
			ssc.close();
		} else if (request.position() > 0) {
			// client request
			System.out.println("Server: read "
					+ new String(request.array(), 0, request.position()));
			ByteBuffer reply = mgr.getAppSendBuffer();
			request.flip();
			reply.put(request);
			request.compact();
			handleWrite(sk);
		}
	}

	// Process OP_WRITE
	void handleWrite(SelectionKey sk) throws IOException {
		SSLEngineManager mgr = (SSLEngineManager) sk.attachment();
		ByteBuffer reply = mgr.getAppSendBuffer();
		System.out.println("Server: writing " + reply);
		int count = 0;
		while (reply.position() > 0) {
			reply.flip();
			count = mgr.write();
			reply.compact();
			if (count == 0)
				break;
		}
		if (reply.position() > 0) {
			// short write:
			// Register for OP_WRITE and come back here when ready
			sk.interestOps(sk.interestOps() | SelectionKey.OP_WRITE);
		} else {
			// Write succeeded, don’t need OP_WRITE any more
			sk.interestOps(sk.interestOps() & ~SelectionKey.OP_WRITE);
		}
	}

	// The main program follows.
	// Main program
	public static void main(String[] args) throws Exception {
		// System.setProperty("javax.net.debug","ssl");
		// TODO adjust these values to suit your local system.
		// These values are for the JDK SSL samples ‘testkeys’.
		System.setProperty("javax.net.ssl.keyStore", "testkeys");
		System.setProperty("javax.net.ssl.keyStorePassword", "passphrase");
		new SSLNIOEchoServer().run();
		System.out.println("Exiting.");
	}
}